package cn.sskxyz.mongodb.multi;

import java.util.Map;

import org.springframework.boot.autoconfigure.mongo.MongoProperties;
import org.springframework.dao.DataAccessException;
import org.springframework.dao.support.PersistenceExceptionTranslator;
import org.springframework.data.mongodb.MongoDatabaseFactory;

import com.mongodb.ClientSessionOptions;
import com.mongodb.client.ClientSession;
import com.mongodb.client.MongoDatabase;

public class MultiGridFsDatabaseFactory implements MongoDatabaseFactory {

	private Map<String, GridFsMongoDatabaseFactory> gridFsDatabaseFactoryMap;

	private String primary;

	public MultiGridFsDatabaseFactory(Map<String, GridFsMongoDatabaseFactory> gridFsDatabaseFactoryMap,
			String primary) {
		this.gridFsDatabaseFactoryMap = gridFsDatabaseFactoryMap;
		this.primary = primary;
	}

	public MongoProperties getMongoProperties() {
		MongoLinkContext context = MongoDbLinkHolder.get();
		if (context == null || context.getLink() == null) {
			return this.gridFsDatabaseFactoryMap.get(primary).getProperties();
		}

		return this.gridFsDatabaseFactoryMap.get(context.getLink()).getProperties();
	}

	@Override
	public MongoDatabase getMongoDatabase() throws DataAccessException {
		return this.mongoDatabaseFactory().getMongoDatabase();
	}

	@Override
	public MongoDatabase getMongoDatabase(String dbName) throws DataAccessException {
		return this.mongoDatabaseFactory().getMongoDatabase(dbName);
	}

	@Override
	public PersistenceExceptionTranslator getExceptionTranslator() {
		return this.mongoDatabaseFactory().getExceptionTranslator();
	}

	@Override
	public ClientSession getSession(ClientSessionOptions options) {
		return this.mongoDatabaseFactory().getSession(options);
	}

	@Override
	public MongoDatabaseFactory withSession(ClientSession session) {
		return this.mongoDatabaseFactory().withSession(session);
	}

	private MongoDatabaseFactory mongoDatabaseFactory() {
		MongoLinkContext context = MongoDbLinkHolder.get();
		if (context == null || context.getLink() == null) {
			return this.gridFsDatabaseFactoryMap.get(primary);
		}
		return gridFsDatabaseFactoryMap.get(context.getLink());
	}
}
